import paddle_aux
import paddle
"""
Auto-anchor utils
"""
import random
import numpy as np
import yaml
from tqdm import tqdm
from pathlib import Path
from utils.general import colorstr


def check_anchor_order(m):
    a = m.anchors.prod(-1).view(-1)
    da = a[-1] - a[0]
    ds = m.stride[-1] - m.stride[0]
    if da.sign() != ds.sign():
        print('Reversing anchor order')
        m.anchors[:] = m.anchors.flip(0)


def check_anchors(dataset, model, thr=4.0, imgsz=640):
    prefix = colorstr('autoanchor: ')
    print(f'\n{prefix}Analyzing anchors... ', end='')
    m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1]
    shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
    scale = np.random.uniform(0.9, 1.1, size=(tuple(shapes.shape)[0], 1))
    wh = paddle.to_tensor(data=np.concatenate([(l[:, 3:5] * s) for s, l in
        zip(shapes * scale, dataset.labels)])).astype(dtype='float32')

    def metric(k):
        r = wh[:, None] / k[None]
        x = paddle_aux.min(r, 1.0 / r).min(2)[0]
        best = x.max(1)[0]
        aat = (x > 1.0 / thr).astype(dtype='float32').sum(axis=1).mean()
        bpr = (best > 1.0 / thr).astype(dtype='float32').mean()
        return bpr, aat
    anchors = m.anchors.clone() * m.stride.to(m.anchors.device).view(-1, 1, 1)
    bpr, aat = metric(anchors.cpu().view(-1, 2))
    print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}'
        , end='')
    if bpr < 0.98:
        print('. Attempting to improve anchors, please wait...')
        na = m.anchors.numel() // 2
        try:
            anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr,
                gen=1000, verbose=False)
        except Exception as e:
            print(f'{prefix}ERROR: {e}')
        new_bpr = metric(anchors)[0]
        if new_bpr > bpr:
            anchors = paddle.to_tensor(data=anchors, place=m.anchors.device
                ).astype(dtype=m.anchors.dtype)
            m.anchors[:] = anchors.clone().view_as(other=m.anchors
                ) / m.stride.to(m.anchors.device).view(-1, 1, 1)
            check_anchor_order(m)
            print(
                f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.'
                )
        else:
            print(
                f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.'
                )
    print('')


def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0,
    gen=1000, verbose=True):
    """ Creates kmeans-evolved anchors from training dataset

        Arguments:
            dataset: path to data.yaml, or a loaded dataset
            n: number of anchors
            img_size: image size used for training
            thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
            gen: generations to evolve anchors using genetic algorithm
            verbose: print all results

        Return:
            k: kmeans evolved anchors

        Usage:
            from utils.autoanchor import *; _ = kmean_anchors()
    """
    from scipy.cluster.vq import kmeans
    thr = 1.0 / thr
    prefix = colorstr('autoanchor: ')

    def metric(k, wh):
        r = wh[:, None] / k[None]
        x = paddle_aux.min(r, 1.0 / r).min(2)[0]
        return x, x.max(1)[0]

    def anchor_fitness(k):
        _, best = metric(paddle.to_tensor(data=k, dtype='float32'), wh)
        return (best * (best > thr).astype(dtype='float32')).mean()

    def print_results(k):
        k = k[np.argsort(k.prod(axis=1))]
        x, best = metric(k, wh0)
        bpr, aat = (best > thr).astype(dtype='float32').mean(), (x > thr
            ).astype(dtype='float32').mean() * n
        print(
            f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr'
            )
        print(
            f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, past_thr={x[x > thr].mean():.3f}-mean: '
            , end='')
        for i, x in enumerate(k):
            print('%i,%i' % (round(x[0]), round(x[1])), end=',  ' if i < 
                len(k) - 1 else '\n')
        return k
    if isinstance(dataset, str):
        with open(dataset, errors='ignore') as f:
            data_dict = yaml.safe_load(f)
            path = Path(data_dict.get('path') or '')
            for k in ('train', 'val', 'test'):
                if data_dict.get(k):
                    data_dict[k] = str(path / data_dict[k]) if isinstance(
                        data_dict[k], str) else [str(path / x) for x in
                        data_dict[k]]
        from utils.datasets import LoadImagesAndLabels
        dataset = LoadImagesAndLabels(data_dict['train'], augment=True,
            rect=True)
    shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
    wh0 = np.concatenate([(l[:, 3:5] * s) for s, l in zip(shapes, dataset.
        labels)])
    i = (wh0 < 3.0).astype('bool').any(axis=1).sum()
    if i:
        print(
            f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.'
            )
    wh = wh0[(wh0 >= 2.0).astype('bool').any(axis=1)]
    print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...')
    s = wh.std(axis=0)
    k, dist = kmeans(wh / s, n, iter=30)
    assert len(k
        ) == n, f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}'
    k *= s
    wh = paddle.to_tensor(data=wh, dtype='float32')
    wh0 = paddle.to_tensor(data=wh0, dtype='float32')
    k = print_results(k)
    npr = np.random
    f, sh, mp, s = anchor_fitness(k), tuple(k.shape), 0.9, 0.1
    pbar = tqdm(range(gen), desc=
        f'{prefix}Evolving anchors with Genetic Algorithm:')
    for _ in pbar:
        v = np.ones(sh)
        while (v == 1).astype('bool').all():
            v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) *
                s + 1).clip(0.3, 3.0)
        kg = (k.copy() * v).clip(min=2.0)
        fg = anchor_fitness(kg)
        if fg > f:
            f, k = fg, kg.copy()
            pbar.desc = (
                f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
                )
            if verbose:
                print_results(k)
    return print_results(k)
